# -*- coding: utf-8 -*-
"""
Created on Wed Mar 12 14:43:51 2014

@author: Maxim
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve


def get_sine_distribution(nPts):
    """
    Returns sine distribution between 0 and 1 with given number of points
    """
    xx = np.linspace(0.0,1.0,nPts)
    cos_curve = lambda x: 1.0-np.cos(x*np.pi/2.0)
    return np.array([cos_curve(_x) for _x in xx])

def spherical(R):
    pts = np.zeros([2,3])
    pts[1] = np.array([R,0,R])
    return pts,R

def conical(L,R):
    pts = np.zeros([2,3])
    pts[1] = np.array([L, 0.0, R])
    return pts

def conical_spherically_blunted(L,R,Rnose):
    pts = np.zeros([3,3])
    xt = L*L/R*(Rnose*Rnose/(R*R+L*L))**0.5
    yt = xt*R/L
    x0 = xt + (Rnose*Rnose-yt*yt)**0.5
    xa = x0-Rnose
    pts[0] = np.array([xa, 0, 0])
    pts[1] = np.array([xt, 0, yt])
    pts[2] = np.array([L, 0, R])
    return pts

def tangent_ogive(L,R):
    assert L>=R
    pts = np.zeros([2,3])
    pts[1] = np.array([L,0,R])
    rho = (R*R+L*L)/(2.*R)
    return pts, rho

def tangent_ogive_spherically_blunted(L,R,Rnose):
    assert L>=R
    pts = np.zeros([3,3])
    rho = (R*R+L*L)/(2.*R)
    x0 = L-((rho-Rnose)**2.-(rho-R)**2.)**0.5
    yt = Rnose*(rho-R)/(rho-Rnose)
    xt = x0 - (Rnose*Rnose-yt*yt)**0.5
    xa = x0-Rnose
    pts[0] = np.array([xa, 0, 0])
    pts[1] = np.array([xt, 0, yt])
    pts[2] = np.array([L, 0, R])
    return pts, rho

def parabolic(L,R,K,nPts=20):
    K = float(K)
#    assert 0<=K<=1.0
    pts = np.zeros([nPts,3])
    x = np.linspace(0,1,nPts)
    pts[:,0] = x*L
    pts[:,2] = R*(2.*x-K*x*x)/(2.-K)
    return pts

def power_series(L,R,n,nPts=20):
    n = float(n)
#    assert 0.25<=n<=1.0
    pts = np.zeros([nPts,3])
    x = get_sine_distribution(nPts)
    #np.linspace(0,1,nPts)
    pts[:,0] = x*L
    pts[:,2] = R*x**n
    return pts

def haack_series(L,R,C,nPts=20):
    assert 0<=C<=1.5
    pts = np.zeros([nPts,3])
    x = get_sine_distribution(nPts)
    theta = np.arccos(1.0-2*x)
    y = R/(np.pi**0.5)* (theta-np.sin(2*theta)/2+C*np.sin(theta)**3)**0.5
    pts[:,0] = x*L
    pts[:,2] = y
    return pts

def conical_spherically_blunted2(L,R,Rnose):
    f = lambda x: x-conical_spherically_blunted(x,R,Rnose)[0,0]-L
    Lnew = fsolve(f,L)
    ptsNew = conical_spherically_blunted(Lnew,R,Rnose)
    ptsNew[:,0] = ptsNew[:,0]-ptsNew[0,0]
    return ptsNew

def tangent_ogive_spherically_blunted2(L,R,Rnose):
    f = lambda x: x-tangent_ogive_spherically_blunted(x,R,Rnose)[0][0,0]-L
    Lnew = fsolve(f,L)
    ptsNew,rho = tangent_ogive_spherically_blunted(Lnew,R,Rnose)
    ptsNew[:,0] = ptsNew[:,0]-ptsNew[0,0]
    return ptsNew,rho

def conical_flat(L,R,Rnose):
    pts = np.zeros([3,3])
    pts[1,2] = Rnose
    pts[2,0] = L
    pts[2,2] = R
    return pts

def tangent_ogive_flat(L,R,Rnose):
    pts,rho = tangent_ogive(L,R-Rnose)
    pts[:,2] += Rnose
    pts = np.vstack([np.zeros(3),pts])
    return pts,rho

def plot_shape():
    L = 4.0
    D = 2.0
    R = D/2
    Rnose = 0.5
    pts1 = parabolic(L,R,0.9)
    pts2 = conical(L,R)
    pts3 = conical_spherically_blunted(L,R,Rnose)
    pts4,rho1 = tangent_ogive(L,R)
    pts5,rho2 = tangent_ogive_spherically_blunted(L,R,Rnose)
    pts6 = power_series(L,R,0.5)
    pts7 = haack_series(L,R,1.333)
    pts8 = conical_spherically_blunted2(L,R,Rnose)
    plt.figure(1)
    plt.hold(True)
    plt.grid(True)
    plt.plot(pts1[:,0],pts1[:,2])
    plt.plot(pts2[:,0],pts2[:,2])
    plt.plot(pts3[:,0],pts3[:,2])
    plt.plot(pts4[:,0],pts4[:,2])
    plt.plot(pts5[:,0],pts5[:,2])
    plt.plot(pts6[:,0],pts6[:,2],'o')
    plt.plot(pts7[:,0],pts7[:,2],'k*-')
    plt.plot(pts8[:,0],pts8[:,2],'k+-')
    plt.axis('equal')
    plt.show()

def run_test1():
    L = 4.0
    D = 2.0
    R = D/2
    Rnose = 0.5
    pts1 = conical_flat(L,R,Rnose)
    pts2,rho = tangent_ogive_flat(L,R,Rnose)
    plt.figure(1)
    plt.hold(True)
    plt.grid(True)
    plt.plot(pts1[:,0],pts1[:,2],'rs-')
    plt.plot(pts2[:,0],pts2[:,2],'bo-')
    plt.axis('equal')
    plt.show()

def plot_power_series():
    L = 4.0
    D = 4.0
    R = D/2
    K = np.array([0, 0.2, 0.4, 0.6, 0.8, 1.0])
    n = np.array([0, 0.2, 0.4, 0.6, 0.8, 1.0])
    C = np.array([0, 0.3, 0.6, 0.9, 1.2, 1.5])
    
    fig1 = plt.figure(1)
    fig2 = plt.figure(2)
    fig3 = plt.figure(3)
    ax1 = fig1.add_subplot(111)
    ax2 = fig2.add_subplot(111)
    ax3 = fig3.add_subplot(111)
    ax1.hold(True)
    ax2.hold(True)
    ax3.hold(True)
    ax1.axis('equal')
    ax2.axis('equal')
    ax3.axis('equal')
    label1, label2, label3 = list(), list(), list()
    for _k, _n, _c in zip(K,n,C):
        pts1 = parabolic(L,R,_k)
        pts2 = power_series(L,R,_n)
        pts3 = haack_series(L,R,_c)
        ax1.plot(pts1[:,0],pts1[:,2])
        ax2.plot(pts2[:,0],pts2[:,2])
        ax3.plot(pts3[:,0],pts3[:,2])
        label1.append('K=%.1f'%_k)
        label2.append('n=%.1f'%_n)
        label3.append('C=%.1f'%_c)
    ax1.legend(label1,'lower right')
    ax2.legend(label2,'lower right')
    ax3.legend(label3,'lower right')
    ax1.set_title('Parabolic')
    ax2.set_title('Power series')
    ax3.set_title('Haack series')
    plt.show()

    
if __name__=="__main__":
    plot_shape()